import random

import torch.utils.data as data
import torchvision
import torchvision.transforms as transforms

class CIFAR10_Dataset(data.Dataset):
    def __init__(self, dataset, noise_rate, transform=None, target_transform=None):
        self.dataset = list(dataset)
        self.noise_rate = noise_rate
        self.transform = transform
        self.target_transform = target_transform

        # construct label noise (img, clean label, noise label)
        for i in range(len(self.dataset)):
            self.dataset[i] = list(self.dataset[i])
            if random.randint(0, 99) < noise_rate * 100:
                while True:
                    x = random.randint(0, 9)
                    if x != self.dataset[i][1]: 
                        self.dataset[i].append(x)
                        break
            else:
                self.dataset[i].append(self.dataset[i][1])
                
    def __getitem__(self, index):
        img, clean, noise = self.dataset[index]
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            clean = self.target_transform(clean)

        if self.target_transform is not None:
            noise = self.target_transform(noise)

        return img, clean, noise

    def __len__(self):
        return len(self.dataset)

class CIFAR10_Val(data.Dataset):
    def __init__(self, dataset, sample_rate, transform=None, target_transform=None):
        self.dataset = list(dataset)
        self.transform = transform
        self.target_transform = target_transform
        random.shuffle(self.dataset)
        self.dataset = self.dataset[:int(len(dataset) * sample_rate)]

    def __getitem__(self, index):
        img, label = self.dataset[index]
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            label = self.target_transform(label)
        return img, label

    def __len__(self):
        return len(self.dataset)
    
        
if __name__ == '__main__':
    # Image preprocessing modules
    transform = transforms.Compose([
    transforms.Pad(4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32),
    transforms.ToTensor()])

    # CIFAR-10 dataset
    train_dataset = torchvision.datasets.CIFAR10(root='data/',
                                                train=True, 
                                                download=True)

    noise_dataset = CIFAR10_Dataset(train_dataset, 0.0, transform)
    val_dataset = CIFAR10_Val(train_dataset, 0.01, transform)

    cnt = 0
    dic1 = {}
    dic2 = {}
    dic3 = {}

    for i in range(len(val_dataset)):
        dic3[val_dataset[i][1]] = dic3[val_dataset[i][1]] + 1 if dic3.get(val_dataset[i][1]) else 1
    print('len val:', len(val_dataset))
    print('val distribution:')
    print(dic3)

    for i in range(len(train_dataset)):
        dic1[noise_dataset[i][1]] = dic1[noise_dataset[i][1]] + 1 if dic1.get(train_dataset[i][1]) else 1
        dic2[noise_dataset[i][2]] = dic2[noise_dataset[i][2]] + 1 if dic2.get(noise_dataset[i][2]) else 1
        if noise_dataset[i][1] != noise_dataset[i][2]:
            cnt += 1
    print('noise_rate:', cnt / len(train_dataset))
    print('noise distribution:')
    print(dic1)
    print(dic2)
    